{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0、安装依赖包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U pinecone-client qianfan tiktoken langchain pymupdf -i https://pypi.tuna.tsinghua.edu.cn/simple/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1、加载环境变量\n",
    "\n",
    "新建一个env的文件，然后填写如下内容：<br>\n",
    "PINECONE_API_KEY=\"your PINECONE_API_KEY\"<br>\n",
    "QIANFAN_AK=\"your QIANFAN_AK\"<br>\n",
    "QIANFAN_SK=\"your QIANFAN_SK\"\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 通过配置文件加载环境变量\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv(dotenv_path = \"env\",override = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "至此我们就完成了了SDK的鉴权初始化配置"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2、pinecone初始化\n",
    "\n",
    "初始化pinecone，并创建pinecone索引（index）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pinecone\n",
    "from tqdm.autonotebook import tqdm\n",
    "\n",
    "PINECONE_API_KEY = os.getenv(\"PINECONE_API_KEY\")\n",
    "\n",
    "# PINECONE免费用户的默认ENV都是\"gcp-starter\"\n",
    "PINECONE_ENV = \"gcp-starter\"\n",
    "\n",
    "pinecone.init(\n",
    "    api_key=PINECONE_API_KEY,\n",
    "    environment=PINECONE_ENV\n",
    ")\n",
    "\n",
    "index_name = \"qianfan-vdb\"\n",
    "\n",
    "if index_name not in pinecone.list_indexes():\n",
    "    # we create a new index\n",
    "    pinecone.create_index(\n",
    "        name=index_name,\n",
    "        metric='euclidean',\n",
    "        dimension=384,  # 千帆Embedding-V1向量模型是384维，所以这里设置成384\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3、文件切片并保存至向量数据库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import PyMuPDFLoader\n",
    "\n",
    "loader = PyMuPDFLoader(\"./example_data/ai-paper.pdf\")\n",
    "documents = loader.load()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4、切片"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "初始化vectorstore，使用PyMuPDFLoader加载PDF文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "33 documents block loaded\n"
     ]
    }
   ],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=384, chunk_overlap=0, separators=[\"\\n\\n\", \"\\n\", \" \", \"\", \"。\", \"，\"])\n",
    "docs_spilts = text_splitter.split_documents(documents)\n",
    "print(f\"{len(docs_spilts)} documents block loaded\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5、存储"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores import Pinecone\n",
    "from langchain.embeddings import QianfanEmbeddingsEndpoint\n",
    "\n",
    "embeddings = QianfanEmbeddingsEndpoint(model=\"Embedding-V1\")\n",
    "pinecone_vdb = Pinecone.from_documents(docs_spilts, embeddings, index_name=index_name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6、加载已有向量数据库\n",
    "\n",
    "如果向量数据库已存在，并且已有数据，可以通过Pinecone.from_existing_index加载已有向量数据库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<langchain.vectorstores.pinecone.Pinecone at 0x153f3b150>"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.vectorstores import Pinecone\n",
    "from langchain.embeddings import QianfanEmbeddingsEndpoint\n",
    "\n",
    "index_name = \"qianfan-vdb\"\n",
    "embeddings = QianfanEmbeddingsEndpoint(model=\"Embedding-V1\")\n",
    "\n",
    "Pinecone.from_existing_index(index_name=index_name, embedding=embeddings)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7、新建Retriever\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = pinecone_vdb.as_retriever(search_type=\"similarity_score_threshold\", search_kwargs={'score_threshold': 0.5, \"k\":3 })"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8、查看相似度文档"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义query\n",
    "query = \"ERNIE1.0是什么时候发布的？\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/xuwei/opt/anaconda3/lib/python3.11/site-packages/langchain_core/vectorstores.py:313: UserWarning: Relevance scores must be between 0 and 1, got [(Document(page_content='勇在阿里云峰会上正式发布了大语言模型工具通义\\n千问，并宣布此语言模型会接入阿里旗下的所有业\\n务中。\\n百度 ERNIE\\n百 度 于 2019 年 3 月 发 布 预 训 练 模 型\\nERNIE1.0，2019 年 7 月 发 布 ERNIE2.0，2021\\n年 5 月开源四大预训练模型，包括多粒度语言知\\n识模型 ERNIE-Gram、超长文本双向建模预训练\\n模型 ERNIE-Doc、融合场景图知识的跨模态预训\\n练模型 ERNIE-ViL 和语言与视觉一体的预训练\\n模型 ERNIE-UNIMO，2021 年 12 月发布多语言\\n预训练模型 ERNIE-M。百度持续投入大模型的\\n技术创新与产业应用，布局了 NLP、CV、跨模态\\n等大模型，率先提出行业大模型，构建大模型工具\\n与平台，探索产品与社区，在企业端和用户端均有', metadata={'author': 'CNKI', 'creationDate': \"D:20230718144804-08'00'\", 'creator': 'ReaderEx_DIS 2.3.0 Build 4031', 'file_path': './example_data/ai-paper.pdf', 'format': 'PDF 1.6', 'keywords': '', 'modDate': '', 'page': 5.0, 'producer': 'TTKN', 'source': './example_data/ai-paper.pdf', 'subject': '', 'title': '', 'total_pages': 7.0, 'trapped': ''}), 1.07225263), (Document(page_content='勇在阿里云峰会上正式发布了大语言模型工具通义\\n千问，并宣布此语言模型会接入阿里旗下的所有业\\n务中。\\n百度 ERNIE\\n百 度 于 2019 年 3 月 发 布 预 训 练 模 型\\nERNIE1.0，2019 年 7 月 发 布 ERNIE2.0，2021\\n年 5 月开源四大预训练模型，包括多粒度语言知\\n识模型 ERNIE-Gram、超长文本双向建模预训练\\n模型 ERNIE-Doc、融合场景图知识的跨模态预训\\n练模型 ERNIE-ViL 和语言与视觉一体的预训练\\n模型 ERNIE-UNIMO，2021 年 12 月发布多语言\\n预训练模型 ERNIE-M。百度持续投入大模型的\\n技术创新与产业应用，布局了 NLP、CV、跨模态\\n等大模型，率先提出行业大模型，构建大模型工具\\n与平台，探索产品与社区，在企业端和用户端均有', metadata={'author': 'CNKI', 'creationDate': \"D:20230718144804-08'00'\", 'creator': 'ReaderEx_DIS 2.3.0 Build 4031', 'file_path': './example_data/ai-paper.pdf', 'format': 'PDF 1.6', 'keywords': '', 'modDate': '', 'page': 5.0, 'producer': 'TTKN', 'source': './example_data/ai-paper.pdf', 'subject': '', 'title': '', 'total_pages': 7.0, 'trapped': ''}), 1.07225263), (Document(page_content='勇在阿里云峰会上正式发布了大语言模型工具通义\\n千问，并宣布此语言模型会接入阿里旗下的所有业\\n务中。\\n百度 ERNIE\\n百 度 于 2019 年 3 月 发 布 预 训 练 模 型\\nERNIE1.0，2019 年 7 月 发 布 ERNIE2.0，2021\\n年 5 月开源四大预训练模型，包括多粒度语言知\\n识模型 ERNIE-Gram、超长文本双向建模预训练\\n模型 ERNIE-Doc、融合场景图知识的跨模态预训\\n练模型 ERNIE-ViL 和语言与视觉一体的预训练\\n模型 ERNIE-UNIMO，2021 年 12 月发布多语言\\n预训练模型 ERNIE-M。百度持续投入大模型的\\n技术创新与产业应用，布局了 NLP、CV、跨模态\\n等大模型，率先提出行业大模型，构建大模型工具\\n与平台，探索产品与社区，在企业端和用户端均有', metadata={'author': 'CNKI', 'creationDate': \"D:20230718144804-08'00'\", 'creator': 'ReaderEx_DIS 2.3.0 Build 4031', 'file_path': './example_data/ai-paper.pdf', 'format': 'PDF 1.6', 'keywords': '', 'modDate': '', 'page': 5.0, 'producer': 'TTKN', 'source': './example_data/ai-paper.pdf', 'subject': '', 'title': '', 'total_pages': 7.0, 'trapped': ''}), 1.07225263)]\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# 检索与query相似度最高的文档\n",
    "matched_docs = retriever.get_relevant_documents(query)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9、问答"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 新建LLM类对象\n",
    "from langchain.chat_models import QianfanChatEndpoint\n",
    "\n",
    "llm = QianfanChatEndpoint(model=\"ERNIE-Bot-4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQA\n",
    "from langchain.prompts import PromptTemplate\n",
    "\n",
    "CUSTOM_PROMPT = \"\"\"\n",
    "现在你是一个阅读理解机器人，你会阅读并深度理解我给你的文本内容并据此回答我所提出的问题。注意，我给出的问题是：{question} 你需要阅读理解的文本是：{context}\n",
    "\"\"\"\n",
    "\n",
    "qa_chain = RetrievalQA.from_chain_type(llm, retriever=retriever, chain_type_kwargs={\"prompt\": PromptTemplate.from_template(CUSTOM_PROMPT)}, return_source_documents=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/xuwei/opt/anaconda3/lib/python3.11/site-packages/langchain_core/vectorstores.py:313: UserWarning: Relevance scores must be between 0 and 1, got [(Document(page_content='勇在阿里云峰会上正式发布了大语言模型工具通义\\n千问，并宣布此语言模型会接入阿里旗下的所有业\\n务中。\\n百度 ERNIE\\n百 度 于 2019 年 3 月 发 布 预 训 练 模 型\\nERNIE1.0，2019 年 7 月 发 布 ERNIE2.0，2021\\n年 5 月开源四大预训练模型，包括多粒度语言知\\n识模型 ERNIE-Gram、超长文本双向建模预训练\\n模型 ERNIE-Doc、融合场景图知识的跨模态预训\\n练模型 ERNIE-ViL 和语言与视觉一体的预训练\\n模型 ERNIE-UNIMO，2021 年 12 月发布多语言\\n预训练模型 ERNIE-M。百度持续投入大模型的\\n技术创新与产业应用，布局了 NLP、CV、跨模态\\n等大模型，率先提出行业大模型，构建大模型工具\\n与平台，探索产品与社区，在企业端和用户端均有', metadata={'author': 'CNKI', 'creationDate': \"D:20230718144804-08'00'\", 'creator': 'ReaderEx_DIS 2.3.0 Build 4031', 'file_path': './example_data/ai-paper.pdf', 'format': 'PDF 1.6', 'keywords': '', 'modDate': '', 'page': 5.0, 'producer': 'TTKN', 'source': './example_data/ai-paper.pdf', 'subject': '', 'title': '', 'total_pages': 7.0, 'trapped': ''}), 1.07225263), (Document(page_content='勇在阿里云峰会上正式发布了大语言模型工具通义\\n千问，并宣布此语言模型会接入阿里旗下的所有业\\n务中。\\n百度 ERNIE\\n百 度 于 2019 年 3 月 发 布 预 训 练 模 型\\nERNIE1.0，2019 年 7 月 发 布 ERNIE2.0，2021\\n年 5 月开源四大预训练模型，包括多粒度语言知\\n识模型 ERNIE-Gram、超长文本双向建模预训练\\n模型 ERNIE-Doc、融合场景图知识的跨模态预训\\n练模型 ERNIE-ViL 和语言与视觉一体的预训练\\n模型 ERNIE-UNIMO，2021 年 12 月发布多语言\\n预训练模型 ERNIE-M。百度持续投入大模型的\\n技术创新与产业应用，布局了 NLP、CV、跨模态\\n等大模型，率先提出行业大模型，构建大模型工具\\n与平台，探索产品与社区，在企业端和用户端均有', metadata={'author': 'CNKI', 'creationDate': \"D:20230718144804-08'00'\", 'creator': 'ReaderEx_DIS 2.3.0 Build 4031', 'file_path': './example_data/ai-paper.pdf', 'format': 'PDF 1.6', 'keywords': '', 'modDate': '', 'page': 5.0, 'producer': 'TTKN', 'source': './example_data/ai-paper.pdf', 'subject': '', 'title': '', 'total_pages': 7.0, 'trapped': ''}), 1.07225263), (Document(page_content='勇在阿里云峰会上正式发布了大语言模型工具通义\\n千问，并宣布此语言模型会接入阿里旗下的所有业\\n务中。\\n百度 ERNIE\\n百 度 于 2019 年 3 月 发 布 预 训 练 模 型\\nERNIE1.0，2019 年 7 月 发 布 ERNIE2.0，2021\\n年 5 月开源四大预训练模型，包括多粒度语言知\\n识模型 ERNIE-Gram、超长文本双向建模预训练\\n模型 ERNIE-Doc、融合场景图知识的跨模态预训\\n练模型 ERNIE-ViL 和语言与视觉一体的预训练\\n模型 ERNIE-UNIMO，2021 年 12 月发布多语言\\n预训练模型 ERNIE-M。百度持续投入大模型的\\n技术创新与产业应用，布局了 NLP、CV、跨模态\\n等大模型，率先提出行业大模型，构建大模型工具\\n与平台，探索产品与社区，在企业端和用户端均有', metadata={'author': 'CNKI', 'creationDate': \"D:20230718144804-08'00'\", 'creator': 'ReaderEx_DIS 2.3.0 Build 4031', 'file_path': './example_data/ai-paper.pdf', 'format': 'PDF 1.6', 'keywords': '', 'modDate': '', 'page': 5.0, 'producer': 'TTKN', 'source': './example_data/ai-paper.pdf', 'subject': '', 'title': '', 'total_pages': 7.0, 'trapped': ''}), 1.07225263)]\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'ERNIE1.0是什么时候发布的？',\n",
       " 'result': '根据提供的文本内容，ERNIE1.0是在2019年3月发布的。\\n\\n文中提到，百度于2019年3月发布预训练模型ERNIE1.0，2019年7月发布ERNIE2.0。因此，答案是ERNIE1.0在2019年3月发布。',\n",
       " 'source_documents': [Document(page_content='勇在阿里云峰会上正式发布了大语言模型工具通义\\n千问，并宣布此语言模型会接入阿里旗下的所有业\\n务中。\\n百度 ERNIE\\n百 度 于 2019 年 3 月 发 布 预 训 练 模 型\\nERNIE1.0，2019 年 7 月 发 布 ERNIE2.0，2021\\n年 5 月开源四大预训练模型，包括多粒度语言知\\n识模型 ERNIE-Gram、超长文本双向建模预训练\\n模型 ERNIE-Doc、融合场景图知识的跨模态预训\\n练模型 ERNIE-ViL 和语言与视觉一体的预训练\\n模型 ERNIE-UNIMO，2021 年 12 月发布多语言\\n预训练模型 ERNIE-M。百度持续投入大模型的\\n技术创新与产业应用，布局了 NLP、CV、跨模态\\n等大模型，率先提出行业大模型，构建大模型工具\\n与平台，探索产品与社区，在企业端和用户端均有', metadata={'author': 'CNKI', 'creationDate': \"D:20230718144804-08'00'\", 'creator': 'ReaderEx_DIS 2.3.0 Build 4031', 'file_path': './example_data/ai-paper.pdf', 'format': 'PDF 1.6', 'keywords': '', 'modDate': '', 'page': 5.0, 'producer': 'TTKN', 'source': './example_data/ai-paper.pdf', 'subject': '', 'title': '', 'total_pages': 7.0, 'trapped': ''}),\n",
       "  Document(page_content='勇在阿里云峰会上正式发布了大语言模型工具通义\\n千问，并宣布此语言模型会接入阿里旗下的所有业\\n务中。\\n百度 ERNIE\\n百 度 于 2019 年 3 月 发 布 预 训 练 模 型\\nERNIE1.0，2019 年 7 月 发 布 ERNIE2.0，2021\\n年 5 月开源四大预训练模型，包括多粒度语言知\\n识模型 ERNIE-Gram、超长文本双向建模预训练\\n模型 ERNIE-Doc、融合场景图知识的跨模态预训\\n练模型 ERNIE-ViL 和语言与视觉一体的预训练\\n模型 ERNIE-UNIMO，2021 年 12 月发布多语言\\n预训练模型 ERNIE-M。百度持续投入大模型的\\n技术创新与产业应用，布局了 NLP、CV、跨模态\\n等大模型，率先提出行业大模型，构建大模型工具\\n与平台，探索产品与社区，在企业端和用户端均有', metadata={'author': 'CNKI', 'creationDate': \"D:20230718144804-08'00'\", 'creator': 'ReaderEx_DIS 2.3.0 Build 4031', 'file_path': './example_data/ai-paper.pdf', 'format': 'PDF 1.6', 'keywords': '', 'modDate': '', 'page': 5.0, 'producer': 'TTKN', 'source': './example_data/ai-paper.pdf', 'subject': '', 'title': '', 'total_pages': 7.0, 'trapped': ''}),\n",
       "  Document(page_content='勇在阿里云峰会上正式发布了大语言模型工具通义\\n千问，并宣布此语言模型会接入阿里旗下的所有业\\n务中。\\n百度 ERNIE\\n百 度 于 2019 年 3 月 发 布 预 训 练 模 型\\nERNIE1.0，2019 年 7 月 发 布 ERNIE2.0，2021\\n年 5 月开源四大预训练模型，包括多粒度语言知\\n识模型 ERNIE-Gram、超长文本双向建模预训练\\n模型 ERNIE-Doc、融合场景图知识的跨模态预训\\n练模型 ERNIE-ViL 和语言与视觉一体的预训练\\n模型 ERNIE-UNIMO，2021 年 12 月发布多语言\\n预训练模型 ERNIE-M。百度持续投入大模型的\\n技术创新与产业应用，布局了 NLP、CV、跨模态\\n等大模型，率先提出行业大模型，构建大模型工具\\n与平台，探索产品与社区，在企业端和用户端均有', metadata={'author': 'CNKI', 'creationDate': \"D:20230718144804-08'00'\", 'creator': 'ReaderEx_DIS 2.3.0 Build 4031', 'file_path': './example_data/ai-paper.pdf', 'format': 'PDF 1.6', 'keywords': '', 'modDate': '', 'page': 5.0, 'producer': 'TTKN', 'source': './example_data/ai-paper.pdf', 'subject': '', 'title': '', 'total_pages': 7.0, 'trapped': ''})]}"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qa_chain({\"query\": query})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "58f7cb64c3a06383b7f18d2a11305edccbad427293a2b4afa7abe8bfc810d4bb"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
